/*
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "Vector3i.h"
#include "AuxMath.h"
#include <functional>
#include "Prerequisites.h"
#include "Geometry/BoxTests.h"
#include "FreeList.h"

#pragma once

//#define __PROF
//#define __NO_SIMPLIFY

#ifdef __PROF
	static int nSimplifiedNils, nSimplifiedSolids;
#endif

template<class VecType>
class Area
{
public:
	Area() : expo(0), pos(VecType()) {}
	Area(Exponent e, const VecType &p) : expo(e), pos(p) {}
	Exponent expo;
	VecType pos;

	inline Area* subAreas(Area arr[VecType::childCount]) const
	{
		const VecType doubled=pos.doubleVec();
		const int e=expo-1;
		for(unsigned int i=0; i<VecType::childCount; i++)
		{
			arr[i].expo = e;
			arr[i].pos = doubled.positiveNeighbor(i);
		}
		return arr;
	}

	///Precondition: The node is the i-th child of a parent node. Retrieves i.
	__forceinline unsigned int getChildIndex() const
	{
		return getChildIndex(pos);
	}
	__forceinline static unsigned int getChildIndex(const Vector3i &pos)
	{
		VecType diff = pos - pos.half().doubleVec();
		unsigned int childIndex = 0;
		//int bla = diff[0];
		for (unsigned int iDim = 0; iDim < VecType::dimension; iDim++)
			childIndex |= (diff[iDim]<<(VecType::dimension-iDim-1));
		// vAssert(childIndex < (int)VecType::childCount);
		return childIndex;
	}
	inline Box<VecType> toBox() const
	{
		const int boxSize=AuxMath::pow2(expo);
		const VecType vMin=pos*boxSize;
		return Box<VecType>(vMin, vMin + VecType(boxSize-1));
	}
	///Tests whether p is inside the area.
	inline bool contains(const VecType &p) const
	{
		return p.isInsideBox(toBox());
	}
	///Tests whether the area contains otherArea.
	inline bool contains(const Area<VecType>& otherArea) const
	{
		const int iLevelDist=expo-otherArea.expo;
		if(iLevelDist<0)
			return false;
		if(!iLevelDist)
			return (otherArea.pos==pos);
		//other area is smaller
		//since areas can't overlap, we only need to check pos (box.min)
		//here we can make use of integer shifting
		VecType vTest;
		for (unsigned int iDim = 0; iDim < VecType::dimension; iDim++)
			vTest[iDim]=otherArea.pos.getComponent(iDim)>>iLevelDist;
		return (vTest==pos);

	}
};



template<class VecType, class DataType>
class NTree
{
public:
	class Node
	{
	public:
		class SSubNodes
		{
		public:
			Node nodes[VecType::childCount];
		};
		friend class NTree;
	private:
		VoxelType nodeType;
		SSubNodes* subNodes;
		DataType data;
	public:

		Node (const Node &rhs)
		{
			data = rhs.data;
			nodeType = rhs.nodeType;
			if (nodeType == Both)
			{
				//createSubNodes doesn't like nodeType == Both. trick it.
				nodeType = Nil;
				createSubNodes();
				for (unsigned int i = 0; i < VecType::childCount; i++)
					subNodes->nodes[i] = rhs.subNodes->nodes[i];
			}
		}
		Node(const DataType &_data) : nodeType(Solid), subNodes(nullptr), data(_data) {}
		Node(VoxelType _nodeType) : nodeType(_nodeType), subNodes(nullptr)
		{
			if (nodeType == Both)
			{
				nodeType = Nil;
				createSubNodes();
			}
		}
		Node() : nodeType(Nil), subNodes(nullptr) {}

		~Node () {
			//printf("deleting node %lld\n", (long long int)this);
			if (nodeType == Both)
			{
				destroySubNodes();
				nodeType = Nil;
			}
		}

		inline Node& operator = (const Node &rhs)
		{
			data = rhs.data;
			if(nodeType == Both && rhs.nodeType != Both)
			{
				destroySubNodes();
				nodeType = rhs.nodeType;
				return *this;
			}
			if (nodeType != Both && rhs.nodeType == Both)
				createSubNodes();
			nodeType = rhs.nodeType;
			if (nodeType == Both)
			{
				for (unsigned int i = 0; i < VecType::childCount; i++)
					subNodes->nodes[i] = rhs.subNodes->nodes[i];
			}
			//puts("node assigned");
			return *this;
		}

		void createSubNodes()
		{
			vAssert(nodeType != Both);
			subNodes = getPool().checkOut();
			for (unsigned int i = 0; i < VecType::childCount; i++)
			{
				subNodes->nodes[i].nodeType=Nil;
				subNodes->nodes[i].subNodes=nullptr;
				subNodes->nodes[i].data = data;
			}

			nodeType = Both;
		}

		void destroySubNodes()
		{
			vAssert(nodeType == Both);

			for (unsigned int i = 0; i < VecType::childCount; i++)
				if(subNodes->nodes[i].nodeType==Both)
					subNodes->nodes[i].destroySubNodes();

			getPool().checkIn(subNodes);
		}

		Node& getSubNode (unsigned int i) { return subNodes->nodes[i]; }
		const Node& constGetSubNode (const unsigned int& i) const { return subNodes->nodes[i]; }
		const DataType& getData () const { return data; }
		void setData(const DataType& _data) { data = _data; }

		const Node* queryArea(const Area<VecType> &nodeArea, const Area<VecType> &queryArea) const
		{
			if (nodeType == Solid && nodeArea.contains(queryArea)) return this;
			if (nodeArea.expo == queryArea.expo && nodeArea.pos == queryArea.pos) return this;
			if (nodeType == Both)
			{
				Area<VecType> subAreas[VecType::childCount];
				nodeArea.subAreas(subAreas);
				for (unsigned int i = 0; i < VecType::childCount; i++)
				{
					if (subAreas[i].contains(queryArea))
						return subNodes->nodes[i].queryArea(subAreas[i], queryArea);
				}
				std::cout << "Node::queryArea failed!" << std::endl;
				vAssert(false);
			}
			return this;
		}
		
		void insertVoxel(const Area<VecType> &nodeArea, const Area<VecType> &area, DataType voxelData)
		{
			if (nodeArea.expo == area.expo && nodeArea.pos == area.pos)
			{
				if (nodeType == Both)
					destroySubNodes();
				nodeType = Solid;		//Overwrite node
				data = voxelData;
			}
			else
			{
				if (nodeType != Both)
					createSubNodes();

				Area<VecType> subAreas[VecType::childCount];
				nodeArea.subAreas(subAreas);
				for (unsigned int i = 0; i < VecType::childCount; i++)
				{
					if (subAreas[i].contains(area))
					{
						subNodes->nodes[i].insertVoxel(subAreas[i], area, voxelData);
						return;
					}
				}
				std::cout << "Node::insertVoxel failed!" << std::endl;
				vAssert(false);
			}
		}

		void
		simplify(bool bRecurse=true)
		{
			if(nodeType==Both)
			{
				unsigned int nNils = 0;
				unsigned int nSolids = 0;
				DataType commData=subNodes->nodes[0].data;
				for (unsigned int i = 0; i < VecType::childCount; i++)
				{
					Node* pNode=subNodes->nodes + i;
					if(bRecurse)
						pNode->simplify();
					switch(pNode->nodeType)
					{
					case Nil:
						nNils++;
						break;
					case Solid:
						if(pNode->data == commData)
							nSolids++;
						break;
					case Both:
						break;//there was an un-simplifiable child
					}
				}
				if(nNils == VecType::childCount)
				{
#ifdef __PROF
					nSimplifiedNils++;
#endif
					destroySubNodes();
					nodeType=Nil;
					return;
				}
				if(nSolids == VecType::childCount)
				{
#ifdef __PROF
					nSimplifiedSolids++;
#endif
					destroySubNodes();
					nodeType=Solid;
					data=commData;
				}
			}
		}

		inline VoxelType getType() const
		{
			return nodeType;
		}
	};

	friend class Node;

	NTree() : size(0) {}
	NTree(int iLevel) : size(iLevel) {}

	template<class Con>
	NTree(Con& c, const DataType& t) : size(0)
	{
		construct(c, t);
	}

	template<class Con> void
	construct(Con& c, const DataType& t)
	{
		if(size)
			return;
		size=c.getSize();
		Area<VecType> subs[VecType::childCount];
		subAreas(subs);
#ifdef __PROF
		nSimplifiedNils=0;
		nSimplifiedSolids=0;
#endif
		for(unsigned int i=0; i<VecType::childCount; i++)
			scanNNode(roots[i], subs[i], t, c);

		lower();
	}

	~NTree()
	{
		//printf("destroying NTree %lld\n", (long long int)this);
	}

	Node& getRoot(unsigned int i) { return roots[i]; }
	const Node& constGetRoot(const unsigned int& i) const { return roots[i]; }
	const Exponent& getSize() const { return size; }

	template<class Con>
	static void
	scanNNode(Node& node, const Area<VecType>& scanArea, const DataType& t, Con& c)
	{
		const VoxelType nodeType=c.fn(scanArea.toBox());
		if(scanArea.expo<=0)
		{
			node.nodeType = nodeType;
			node.data = t;
			if (nodeType == Both) node.nodeType = Nil;
			return;
		}
		//printf("checking area %s, %s", box.min.toString().c_str(), box.max.toString().c_str());
		switch(nodeType)
		{
		case Nil:
			return;
		case Solid:
			node.nodeType=Solid;
			node.data=t;
			return;
		default:
			Area<VecType> subs[VecType::childCount];
			scanArea.subAreas(subs);
			node.createSubNodes();
			for(unsigned int i=0; i<VecType::childCount; i++)
				scanNNode(node.subNodes->nodes[i], subs[i], t, c);
#ifndef __NO_SIMPLIFY
			node.simplify(false);
#endif
			break;
		}
	}

	const Node* queryArea(const Area<VecType> &queryArea) const
	{
		VecType neighbors[VecType::childCount];
		VecType pos(-1);
		pos.positiveNeighbors(neighbors);
		for (unsigned int i = 0; i < VecType::childCount; i++) {
			Area<VecType> subArea(size, neighbors[i]);
			if (subArea.contains(queryArea))
				return roots[i].queryArea(subArea, queryArea);
		}
		return nullptr;
	}

	const VoxelType queryAreaType(const Area<VecType> area) const
	{
		const Node* pNode=queryArea(area);
		if(pNode==nullptr)
			return Nil;
		return pNode->getType();
	}

	void insertVoxel(const Area<VecType> &area, DataType voxelData)
	{
		VecType neighbors[VecType::childCount];
		VecType pos(-1);
		pos.positiveNeighbors(neighbors);
		for (unsigned int i = 0; i < VecType::childCount; i++)
		{
			Area<VecType> subArea(size, neighbors[i]);
			if (subArea.contains(area))
			{
				roots[i].insertVoxel(subArea, area, voxelData);
				return;
			}
		}
		std::cout << "Tree::insertVoxel failed (trying to insert voxel at pos " << area.pos.toString() << ")" << std::endl;
		vAssert(false);		//area not found at all
	}

	Area<VecType>*
	subAreas(Area<VecType> area[VecType::childCount]) const
	{
		VecType neighbors[VecType::childCount];
		VecType pos(-1);
		pos.positiveNeighbors(neighbors);
		for (unsigned int i = 0; i < VecType::childCount; i++)
			area[i]=Area<VecType>(size, neighbors[i]);
		return area;
	}

	void
	raise(const unsigned int& iLevel)
	{
		if(size>=iLevel)
			return;//nothing to rise, already bigger

		//temp data
		VoxelType nodeType;
		typename Node::SSubNodes* subNodes;
		DataType data;

		for (unsigned int i = 0; i < VecType::childCount; i++)
		{
			unsigned int iTargetIndex = VecType::childCount-i-1;
			//move root[i] to root[i].subnodes->nodes[iTargetIndex]
			//set other nodes to Nil root[i].subnodes->nodes[j] = Node(Nil)

			nodeType=roots[i].nodeType;
			subNodes=roots[i].subNodes;
			data=roots[i].data;

			//now it is safe to clear roots[i]
			roots[i].nodeType=Nil;
			roots[i].subNodes=nullptr;

			roots[i].createSubNodes();
			//suubnodes are nil by default

			roots[i].subNodes->nodes[iTargetIndex].nodeType=nodeType;
			roots[i].subNodes->nodes[iTargetIndex].subNodes=subNodes;
			roots[i].subNodes->nodes[iTargetIndex].data=data;

		}

		size++;
		raise(iLevel);
	}

	//remove levels if the exponent of the tree could be smaller
	void
	lower()
	{
		//trivial path: roots are all nil
		bool bAllNil=true;
		for (unsigned int iRoot = 0; iRoot < VecType::childCount; iRoot++)
		{
			if(roots[iRoot].getType()!=Nil)
			{
				bAllNil=false;
				break;
			}
		}
		if(bAllNil)
		{
			size=1;
			return;
		}

		//non-trivial path
		for (unsigned int iRoot = 0; iRoot < VecType::childCount; iRoot++)
		{
			Node* pRoot=roots+iRoot;
			if(pRoot->getType()!=Both)
				return;
			unsigned int iInnerChild=VecType::childCount-iRoot-1;

			for(unsigned int iChild=0; iChild<VecType::childCount; iChild++)
				if(iChild!=iInnerChild)
					if(pRoot->constGetSubNode(iChild).getType()!=Nil)
						return;
		}
		//can lower the tree
		//dummy node
		VoxelType nodeType;
		typename Node::SSubNodes* subNodes;
		DataType data;
		for (unsigned int iRoot = 0; iRoot < VecType::childCount; iRoot++)
		{
			const unsigned int iInnerChild=VecType::childCount-iRoot-1;
			Node& root=roots[iRoot];

			//this node will be eliminated
			//copy data
			subNodes=root.subNodes->nodes[iInnerChild].subNodes;
			nodeType=root.subNodes->nodes[iInnerChild].nodeType;
			data=root.subNodes->nodes[iInnerChild].data;
			
			//fake data in root
			root.subNodes->nodes[iInnerChild].nodeType=Nil;
			root.subNodes->nodes[iInnerChild].subNodes=nullptr;
			root.destroySubNodes();
			//copy contents
			root.nodeType=nodeType;
			root.data=data;
			root.subNodes=subNodes;
		}
		size--;
		lower();
	}

	Box<VecType>
	getBox() const
	{
		int treeExt=AuxMath::pow2(size);
		return Box<VecType>(VecType(treeExt)*(-1), VecType(treeExt));
	}

	void
	simplify()
	{
		for (unsigned int i = 0; i < VecType::childCount; i++)
			roots[i].simplify();
	}

private:
	Node roots[VecType::childCount];
	Exponent size;

	typedef typename Node::SSubNodes subNode_type;

	static FreeList<subNode_type>&
	getPool()
	{
		static FreeList<subNode_type> pool;
		return pool;
	}
};
